# server.py
import asyncio
import json
import logging
import websockets
from typing import Dict, Optional
from dataclasses import dataclass

from asr import ASRSession

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# 模拟 JWT 验证（实际应使用 pyjwt 验证）
def verify_token(token: str) -> bool:
    # TODO: 实际项目中使用 jwt.decode 验证签名和过期时间
    return token.startswith("mock_") or len(token) > 10

# 存储活跃连接（可用于广播或中断）
active_connections: Dict[str, 'ClientConnection'] = {}

@dataclass
class ClientConnection:
    websocket: websockets.WebSocketServerProtocol
    client_id: str
    session_id: str
    task: Optional[asyncio.Task] = None
    asr_session: Optional[ASRSession] = None  # 新增

class RealtimeAssistantServer:
    def __init__(self, host='localhost', port=8765):
        self.host = host
        self.port = port
        self.shutdown_event = asyncio.Event()

    async def handle_connection(self, websocket: websockets.WebSocketServerProtocol):
        client_id = f"client-{id(websocket) % 10000:04d}"
        connection = ClientConnection(
            websocket=websocket,
            client_id=client_id,
            session_id=""
        )
        active_connections[client_id] = connection

        logger.info(f"[{client_id}] 客户端连接建立")

        try:
            await self._handle_client(connection)
        except websockets.exceptions.ConnectionClosed as e:
            logger.info(f"[{client_id}] 连接关闭: {e}")
        except Exception as e:
            logger.error(f"[{client_id}] 处理异常: {e}", exc_info=True)
            await self._send_error(websocket, "internal_error", str(e))
        finally:
            if client_id in active_connections:
                del active_connections[client_id]

    async def _handle_client(self, connection: ClientConnection):
        websocket = connection.websocket
        client_id = connection.client_id

        # 等待 session 消息
        try:
            message = await websocket.recv()
            if isinstance(message, str):
                data = json.loads(message)
            else:
                logger.warning(f"[{client_id}] 收到非文本消息作为第一条消息")
                return

            if data.get("type") == "session":
                token = data.get("token")
                session_id = data.get("id")

                # if not token or not verify_token(token):
                if 1 != 1:
                    await self._send_error(websocket, "auth_failed", "无效的 token")
                    await websocket.close(reason="Invalid token")
                    return

                connection.session_id = session_id or f"sess-{client_id}"

                await websocket.send(
                    json.dumps({
                        "type": "connection-established",
                        "message": "连接建立成功，可以开始对话了！"
                    })
                )
                logger.info(f"[{client_id}] 会话建立成功，session_id={connection.session_id}")

            else:
                await self._send_error(websocket, "invalid_protocol", "第一条消息必须是 session")
                await websocket.close(reason="First message must be session")
                return

        except asyncio.TimeoutError:
            await self._send_error(websocket, "timeout", "会话建立超时")
            await websocket.close(reason="Session setup timeout")
            return
        except Exception as e:
            await self._send_error(websocket, "session_failed", f"会话建立失败: {str(e)}")
            return

        # 进入主消息循环
        await self._main_message_loop(connection)

    async def _main_message_loop(self, connection: ClientConnection):
        """主消息处理循环"""
        websocket = connection.websocket
        client_id = connection.client_id

        loop = asyncio.get_running_loop()

        while True:
            try:
                message = await asyncio.wait_for(websocket.recv(), timeout=300)  # 5分钟超时

                if isinstance(message, str):
                    # 文本消息
                    data = json.loads(message)
                    msg_type = data.get("type")

                    if msg_type == "audio":
                        action = data.get("data")
                        if action == "start":
                            logger.info(f"[{client_id}] 收到音频开始，启动 ASR")

                            # 创建 ASR 会话
                            def on_asr_stream(text):
                                future = asyncio.run_coroutine_threadsafe(
                                    websocket.send(json.dumps({
                                        "type": "asr-result-stream",
                                        "text": text
                                    })),
                                    loop
                                )
                                try:
                                    future.result(timeout=1.0)
                                except:
                                    pass

                            def on_asr_end(text):
                                future = asyncio.run_coroutine_threadsafe(
                                    websocket.send(json.dumps({
                                        "type": "asr-result-end",
                                        "text": text
                                    })),
                                    loop
                                )
                                try:
                                    future.result(timeout=1.0)
                                except:
                                    pass

                                # TODO: 触发 Agent 处理（下一步实现）

                            def on_asr_error(msg):
                                asyncio.create_task(self._send_error(websocket, "asr_failed", msg))

                            def on_asr_close():
                                logger.info(f"[{client_id}] ASR 会话关闭")
                                connection.asr_session = None

                            asr_session = ASRSession(
                                on_result=on_asr_stream,
                                on_final=on_asr_end,
                                on_error=on_asr_error,
                                on_close=on_asr_close
                            )
                            asr_session.set_event_loop(asyncio.get_running_loop())
                            asr_session.start()
                            connection.asr_session = asr_session

                            # 发送开始状态
                            await websocket.send(
                                json.dumps({
                                    "type": "asr-result-start",
                                    "status": "recognizing",
                                    "message": "正在识别语音..."
                                })
                            )

                        elif action == "end":
                            logger.info(f"[{client_id}] 收到音频结束，停止 ASR")
                            if connection.asr_session:
                                connection.asr_session.stop()
                    elif msg_type == "text-input":
                        text = data.get("text", "").strip()
                        if text:
                            logger.info(f"[{client_id}] 收到文本输入: {text}")
                            # TODO: 直接触发 Agent（下一步实现）
                    elif msg_type == "interrupt-signal":
                        logger.info(f"[{client_id}] 收到中断信号")
                        # 中断 ASR
                        if connection.asr_session:
                            connection.asr_session.close()
                            connection.asr_session = None

                        await websocket.send(
                            json.dumps({
                                "type": "status-update",
                                "status": "idle",
                                "message": "已中断"
                            })
                        )
                    else:
                        logger.warning(f"[{client_id}] 未知消息类型: {msg_type}")

                elif isinstance(message, bytes):
                    logger.debug(f"[{client_id}] 收到音频数据块，大小: {len(message)} 字节")

                    # 必须已经启动了 ASR 会话
                    if connection.asr_session and connection.asr_session.running:
                        # 直接发送音频帧（假设客户端发送的是 16kHz PCM）
                        connection.asr_session.send_audio(message)
                    else:
                        logger.warning(f"[{client_id}] 收到音频数据但 ASR 未启动或已停止")


            except asyncio.TimeoutError:
                logger.info(f"[{client_id}] 连接超时，关闭")
                await websocket.close(reason="Idle timeout")
                break
            except websockets.exceptions.ConnectionClosed:
                break
            except Exception as e:
                logger.error(f"[{client_id}] 消息处理异常: {e}", exc_info=True)
                await self._send_error(websocket, "message_error", str(e))

    async def _send_error(self, websocket, code: str, message: str):
        """发送错误消息"""
        try:
            await websocket.send(
                json.dumps({
                    "type": "error",
                    "code": code,
                    "message": message
                })
            )
        except:
            pass  # 发送失败也忽略

    async def start(self):
        """启动服务器"""
        async with websockets.serve(
            self.handle_connection,
            self.host,
            self.port,
            ping_interval=20,
            ping_timeout=30,
            max_size=1024 * 1024 * 10  # 10MB
        ):
            logger.info(f"WebSocket 服务已启动: ws://{self.host}:{self.port}")
            await self.shutdown_event.wait()

    async def shutdown(self):
        """关闭服务器"""
        self.shutdown_event.set()
        # 关闭所有连接
        for conn in list(active_connections.values()):
            await conn.websocket.close(reason="Server shutdown")
        logger.info("服务器已关闭")

# ======================
# 主程序入口
# ======================
if __name__ == "__main__":
    server = RealtimeAssistantServer(host='localhost', port=8765)
    try:
        asyncio.run(server.start())
    except KeyboardInterrupt:
        print("\n正在关闭服务器...")
        asyncio.run(server.shutdown())